// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <type_traits>

#include <pollux/common/base/exceptions.h>
#include <pollux/vector/builder_type_utils.h>
#include <pollux/vector/flat_vector.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux {
    /*
     * NOTE - biased vector is stored solely as a standard numeric flat array in
     * 2 buffers.
     *
     * Buffer order: [nullData, valueData]
     */
    template<typename T>
    BiasVector<T>::BiasVector(
        pollux::memory::MemoryPool *pool,
        BufferPtr nulls,
        size_t length,
        TypeKind valueType,
        BufferPtr values,
        T bias,
        const SimpleVectorStats<T> &stats,
        std::optional<vector_size_t> distinctCount,
        std::optional<vector_size_t> nullCount,
        std::optional<bool> sorted,
        std::optional<ByteCount> represented_bytes,
        std::optional<ByteCount> storageByteCount)
        : SimpleVector<T>(
              pool,
              CppToType<T>::create(),
              VectorEncoding::Simple::BIASED,
              nulls,
              length,
              stats,
              distinctCount,
              nullCount,
              sorted,
              represented_bytes,
              storageByteCount),
          valueType_(valueType),
          values_(std::move(values)),
          bias_(bias) {
        POLLUX_CHECK(
            valueType_ == TypeKind::INTEGER || valueType_ == TypeKind::SMALLINT ||
            valueType_ == TypeKind::TINYINT,
            "Invalid array type for biased values");

        biasBuffer_ = simd::setAll(bias_);
        rawValues_ = values_->as<uint8_t>();
        BaseVector::in_memory_bytes_ += values_->size();
    }

    template<typename T>

    std::unique_ptr<SimpleVector<uint64_t> > BiasVector<T>::hash_all() const {
        // TODO T70734527 dealing with zero length vector
        if (BaseVector::length_ == 0) {
            return nullptr;
        }
        // If there is at least one value, then indices_ is set and has a pool.
        BufferPtr hashes =
                AlignedBuffer::allocate<uint64_t>(BaseVector::length_, values_->pool());
        uint64_t *rawHashes = hashes->asMutable<uint64_t>();
        for (size_t i = 0; i < BaseVector::length_; ++i) {
            rawHashes[i] = this->hash_value_at(i);
        }
        return std::make_unique<FlatVector<uint64_t> >(
            BaseVector::pool_,
            BIGINT(),
            BufferPtr(nullptr),
            BaseVector::length_,
            hashes,
            std::vector<BufferPtr>(0) /*stringBuffers*/,
            SimpleVectorStats<uint64_t>{},
            std::nullopt /*distinctValueCount*/,
            0 /* nullCount */,
            false /*isSorted*/,
            sizeof(uint64_t) * BaseVector::length_ /*represented_bytes*/);
    }

    template<typename T>
    const T BiasVector<T>::valueAtFast(vector_size_t idx) const {
        switch (valueType_) {
            case TypeKind::INTEGER:
                return bias_ + reinterpret_cast<const int32_t *>(rawValues_)[idx];
            case TypeKind::SMALLINT:
                return bias_ + reinterpret_cast<const int16_t *>(rawValues_)[idx];
            case TypeKind::TINYINT:
                return bias_ + reinterpret_cast<const int8_t *>(rawValues_)[idx];
            default:
                POLLUX_UNSUPPORTED("Invalid type");
        }
    }

    template<typename T>
    ksimd::batch<T> BiasVector<T>::loadSIMDValueBufferAt(size_t index) const {
        if constexpr (std::is_same_v<T, int64_t>) {
            switch (valueType_) {
                case TypeKind::INTEGER:
                    return biasBuffer_ + loadSIMDInternal<int32_t>(index);
                case TypeKind::SMALLINT:
                    return biasBuffer_ + loadSIMDInternal<int16_t>(index);
                case TypeKind::TINYINT:
                    return biasBuffer_ + loadSIMDInternal<int8_t>(index);
                default:
                    POLLUX_UNSUPPORTED("Invalid type");
            }
        } else if constexpr (std::is_same_v<T, int32_t>) {
            switch (valueType_) {
                case TypeKind::SMALLINT:
                    return biasBuffer_ + loadSIMDInternal<int16_t>(index);
                case TypeKind::TINYINT:
                    return biasBuffer_ + loadSIMDInternal<int8_t>(index);
                default:
                    POLLUX_UNSUPPORTED("Invalid type");
            }
        } else if constexpr (std::is_same_v<T, int16_t>) {
            switch (valueType_) {
                case TypeKind::TINYINT:
                    return biasBuffer_ + loadSIMDInternal<int8_t>(index);
                default:
                    POLLUX_UNSUPPORTED("Invalid type");
            }
        } else {
            POLLUX_UNSUPPORTED("Unsupported type for biased vector");
        }
    }
} // namespace kumo::pollux
